{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.optim as optim\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import numpy as np\n",
    "from torch.autograd import Variable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "download=True\n",
    "batch_size=10\n",
    "\n",
    "transform = transforms.Compose([transforms.ToTensor(),\n",
    "                                transforms.Normalize((0.5, 0.5, 0.5),\n",
    "                                                     (0.5, 0.5, 0.5))])\n",
    "trainset = torchvision.datasets.MNIST(root='../data', train=True,\n",
    "                                          download=download,\n",
    "                                          transform=transform)\n",
    "trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,\n",
    "                                              shuffle=True, num_workers=2)\n",
    "testset = torchvision.datasets.MNIST(root='../data', train=False,\n",
    "                                     download=download,\n",
    "                                     transform=transform)\n",
    "testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,\n",
    "                                         shuffle=False, num_workers=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset MNIST\n",
       "    Number of datapoints: 60000\n",
       "    Root location: ../data\n",
       "    Split: Train\n",
       "    StandardTransform\n",
       "Transform: Compose(\n",
       "               ToTensor()\n",
       "               Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))\n",
       "           )"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset MNIST\n",
       "    Number of datapoints: 10000\n",
       "    Root location: ../data\n",
       "    Split: Test\n",
       "    StandardTransform\n",
       "Transform: Compose(\n",
       "               ToTensor()\n",
       "               Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))\n",
       "           )"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "testset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([28, 28])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainset.data[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ImprovedGAN_Discriminator(nn.Module):\n",
    "\n",
    "    def __init__(self, featmap_dim=512, n_channel=1, use_gpu=False,\n",
    "                 n_B=128, n_C=16):\n",
    "        \"\"\"\n",
    "        Minibatch discrimination: learn a tensor to encode side information\n",
    "        from other examples in the same minibatch.\n",
    "        \"\"\"\n",
    "        super(ImprovedGAN_Discriminator, self).__init__()\n",
    "        self.use_gpu = use_gpu\n",
    "        self.n_B = n_B\n",
    "        self.n_C = n_C\n",
    "        self.featmap_dim = featmap_dim\n",
    "\n",
    "        self.conv1 = nn.Conv2d(n_channel, featmap_dim // 4, 5,\n",
    "                               stride=2, padding=2)\n",
    "\n",
    "        self.conv2 = nn.Conv2d(featmap_dim // 4, featmap_dim // 2, 5,\n",
    "                               stride=2, padding=2)\n",
    "        self.BN2 = nn.BatchNorm2d(featmap_dim // 2)\n",
    "\n",
    "        self.conv3 = nn.Conv2d(featmap_dim // 2, featmap_dim, 5,\n",
    "                               stride=2, padding=2)\n",
    "        self.BN3 = nn.BatchNorm2d(featmap_dim)\n",
    "\n",
    "        T_ten_init = torch.randn(featmap_dim * 4 * 4, n_B * n_C) * 0.1\n",
    "        self.T_tensor = nn.Parameter(T_ten_init, requires_grad=True)\n",
    "        self.fc = nn.Linear(featmap_dim * 4 * 4 + n_B, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        Architecture is similar to DCGANs\n",
    "        Add minibatch discrimination => Improved GAN.\n",
    "        \"\"\"\n",
    "        x = F.leaky_relu(self.conv1(x), negative_slope=0.2)\n",
    "        x = F.leaky_relu(self.BN2(self.conv2(x)), negative_slope=0.2)\n",
    "        x = F.leaky_relu(self.BN3(self.conv3(x)), negative_slope=0.2)\n",
    "        x = x.view(-1, self.featmap_dim * 4 * 4)\n",
    "\n",
    "        # #### Minibatch Discrimination ###\n",
    "        T_tensor = self.T_tensor\n",
    "        if self.use_gpu:\n",
    "            T_tensor = T_tensor.cuda()\n",
    "\n",
    "        Ms = x.mm(T_tensor)\n",
    "        Ms = Ms.view(-1, self.n_B, self.n_C)\n",
    "\n",
    "        out_tensor = []\n",
    "        for i in range(Ms.size()[0]):\n",
    "\n",
    "            out_i = None\n",
    "            for j in range(Ms.size()[0]):\n",
    "                o_i = torch.sum(torch.abs(Ms[i, :, :] - Ms[j, :, :]), 1)\n",
    "                o_i = torch.exp(-o_i)\n",
    "                if out_i is None:\n",
    "                    out_i = o_i\n",
    "                else:\n",
    "                    out_i = out_i + o_i\n",
    "\n",
    "            out_tensor.append(out_i)\n",
    "\n",
    "        out_T = torch.cat(tuple(out_tensor)).view(Ms.size()[0], self.n_B)\n",
    "        x = torch.cat((x, out_T), 1)\n",
    "        # #### Minibatch Discrimination ###\n",
    "\n",
    "        x = F.sigmoid(self.fc(x))\n",
    "\n",
    "        return x"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
